//
// This file is part of cpp-lib, a library of C++ functions and classes.
// Cpp-lib is copyright (C) 2004 and onwards, The CPL Interest Group.
// 
// This library is free software; you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation; either version 2.1 of the License, or (at
// your option) any later version.
// 
// This library is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser
// General Public License for more details.
// 
// You should have received a copy of the GNU Lesser General Public License
// along with this library (see the file COPYING); if not, write to the
// Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
// 02111-1307  USA
//
// $Id$
//


#ifndef CPP_LIB_BANDED_MATRIX_H
#define CPP_LIB_BANDED_MATRIX_H

#include <algorithm>
#include <exception>
#include <stdexcept>

#include <cmath>

#include "cpp-lib/matrix.h"


namespace cpl {

namespace math {


/// Specialization for banded matrix with fixed upper and lower bandwidth.

/// A( i , j ) != 0 implies 
///
///   -A.upper_bandwidth() <= i - j <= A.lower_bandwidth() .
///
/// That is, upper_bandwidth() == 0 means lower triangular, 
/// lower_bandwidth() == 0 means upper triangular, both 0 means diagonal
/// matrix.

template< typename T > struct
matrix< T , banded > {

  /// Initialize members and assert() sensible values.

  matrix( 
    const long rows , 
    const long columns , 
    const long lower_bandwidth ,
    const long upper_bandwidth
  ) 
  : columns_( ( assert( columns >= 1 ) , columns ) ) ,
    lower_bandwidth_( (
      assert( lower_bandwidth >= 0 && lower_bandwidth <= rows - 1 ) ,
      lower_bandwidth 
    ) ) ,
    upper_bandwidth_( (
      assert( upper_bandwidth >= 0 && upper_bandwidth <= columns - 1 ) ,
      upper_bandwidth 
    ) ) ,
    data( ( assert( rows >= 1 ) , rows ) , 
          lower_bandwidth + upper_bandwidth + 1 )
    {}

  /// Default constructor makes empty matrix (zero rows and columns).

  matrix()
  : columns_( 0 ) ,
    lower_bandwidth_( 0 ) ,
    upper_bandwidth_( 0 ) ,
    data() 
  {}


  /// Templated copy constructor.

  /// \todo The members are initialized and immediately overwritten.
 
  template< typename tag >
  matrix( matrix< T , tag > const& A ) 
  : columns_( 0 ) , 
    lower_bandwidth_( 0 ) , 
    upper_bandwidth_( 0 ) , 
    data() 
  { *this = A ; }


  /// Swap with other banded matrix.

  void swap( matrix< T , banded >& ) ;


  /// Set all elements within (structural) nonzero band to specified value.

  void fill( T const& x ) { data.fill( x ) ; }


  //@{
  /// \return Reference to element at row i, column j.
  
  /// The non-const operator() asserts that (i,j) is within the nonzero band!
  /// 

  T& operator()( long i , long j ) ;

  /// \return Constant reference to element at row i, column j.  

  /// If (i,j) is out of the nonzero band, return a reference to zero.

  T const& operator()( long i , long j ) const ;

  //@}

  //@{
  /// \return Number of rows.
  long rows() const { return data.rows() ; }
  /// \return Number of columns.
  long columns() const { return columns_ ; }
  //@}

  long get_upper_bandwidth() const { return upper_bandwidth_ ; }
  long get_lower_bandwidth() const { return lower_bandwidth_ ; }


  /// Templated assignment operator.

  /// Semantics: If dimensions disagree or lower or upper bandwidth of
  /// \a rhs is greater than ours, reconstruct ourselves.
  ///
  /// Else, leave our bandwidth unchanged and assign to each element
  /// within.
  ///

  template< typename tag >
  matrix< T , banded > const& operator=( matrix< T , tag > const& rhs ) ;


private:

  /// Number of columns.
  long columns_ ;

  /// Lower bandwidth.
  long lower_bandwidth_ ;

  /// Upper bandwidth.
  long upper_bandwidth_ ;

  /// Data matrix.  

  /// The banded matrix is stored row-wise in data, i.e. the nonzero
  /// part of each row of the banded matrix is a row of data.

  matrix< T , dense > data ;

} ;


/// Banded matrix with constant subdiagonals, \a n rows.  

/// A has the form
/// \f[
/// \begin{bmatrix}
/// x &        &   \cr
///   & \ddots &   \cr 
///   &        & x
/// \end{bmatrix},
/// \f]
/// where \a x is the vector argument interpreted as a row vector.
/// \a x is repeated \a n times, hence the returned matrix has \a n rows and
/// \a n + \a x.size() - 1 columns.

template< typename T >
matrix< T , banded > constant_subdiagonals( std::vector< T > x , long n ) ;


/// \return The (\a n * \a n) identity matrix.

template< typename T > const matrix< T , banded > identity( long n ) ;


/// Banded Cholesky factorization. 

/// \param A A symmetric positive definite banded matrix.  Symmetry
/// isn't checked, we only use the upper triangle of A.
/// \return A lower triangular banded matrix \a L such that 
/// \a L \a L^T = \a A.
/// \exception std::runtime_error If A isn't positive definite.

template< typename T >
const matrix< T , banded > 
cholesky( matrix< T , banded > const& A ) ;


/// Solve Ax = b by forward- or backsubstitution.

/// \param A A nonsingular lower or upper triangular square matrix.
/// \param b a column vector.
/// \return The solution x.
/// \exception std::runtime_error If A is singular, an exception
/// containing an appropriate message is thrown.

template< typename T , typename tag1 , typename tag2 >
const matrix< T , dense >
triangular_solve( matrix< T , tag1 > const& A , matrix< T , tag2 > const& b ) ;


} // namespace math

} // namespace cpl


//
// Template definitions.
//

template< typename T > inline T&
cpl::math::matrix< T , cpl::math::banded >::operator()( long i , long j ) {

  assert( 1 <= i && i <= rows() ) ;

  assert( -upper_bandwidth_ <= i - j ) ;
  assert( i - j <= lower_bandwidth_ )  ;
  return data( i , j - i + lower_bandwidth_ + 1 ) ;

}


template< typename T > inline T const&
cpl::math::matrix< T , cpl::math::banded >::operator()( long i , long j ) const {

  static const T zero = T() ;

  if( -upper_bandwidth_ > i - j || i - j > lower_bandwidth_ )
    return zero ;

  return data( i , j - i + lower_bandwidth_ + 1 ) ;

}


template< typename T >
void cpl::math::matrix< T , cpl::math::banded >::swap( cpl::math::matrix< T , cpl::math::banded >& other ) {

  std::swap( columns_ , other.columns_ ) ;
  std::swap( lower_bandwidth_ , other.lower_bandwidth_ ) ;
  std::swap( upper_bandwidth_ , other.upper_bandwidth_ ) ;
  data.swap( other.data ) ;

}


template< typename T >
cpl::math::matrix< T , cpl::math::banded > 
cpl::math::constant_subdiagonals( std::vector< T > x , long n ) {

  assert( x.size() ) ;

  matrix< T , cpl::math::banded > ret( n , n + x.size() - 1 , 0 , x.size() - 1 ) ;

  for( long i = 1 ; i <= ret.rows() ; ++i ) 
    for( unsigned long j = 0 ; j < x.size() ; ++j )
      ret( i , i + j ) = x[ j ] ;

  return ret ;

}

template< typename T >
const cpl::math::matrix< T , cpl::math::banded >
cpl::math::cholesky( cpl::math::matrix< T , cpl::math::banded > const& A ) {

  assert( A.rows() == A.columns() ) ;

  // Result is lower triangular with same bandwidth as A.  Cf. Golub,
  // van Loan, Section 4.3.5.
  //
  // This implementation follows the NR chapter 2.9 (Press et al.).
  
  cpl::math::matrix< T , cpl::math::banded > 
    L( A.rows() , A.columns() , upper_bandwidth( A ) , 0 ) ;

  for( long i = 1 ; i <= L.rows() ; ++i ) {

    T sum = T() ;

    for( long k = cpl::math::first_column( L , i ) ; k < i ; ++k ) 
    { sum -= std::pow( L( i , k ) , 2 ) ; }

    
    sum += A( i , i ) ;

    if( sum <= 0 ) {

      throw std::runtime_error( "matrix not positive definite" ) ;

    }

    L( i , i ) = std::sqrt( sum ) ;

    for( long j = i + 1 ; j <= cpl::math::last_row( L , i ) ; ++j ) {

      sum = T() ;

      for( 
        long k = std::max( first_column( L , i ) , first_column( L , j ) ) ;
        k < i ;
        ++k
      ) { sum -= L( i , k ) * L( j , k ) ; }

      sum += A( i , j ) ;
      L( j , i ) = sum / L( i , i ) ;

    }

  }

  return L ;

}


template< typename T >
const cpl::math::matrix< T , cpl::math::banded >
cpl::math::identity( long n ) {

  cpl::math::matrix< T , cpl::math::banded > I( n , n , 0 , 0 ) ;

  for( long i = 1 ; i <= n ; ++i ) { I( i , i ) = 1 ; }

  return I ;

}


template< typename T , typename tag1 , typename tag2 >
const cpl::math::matrix< T , cpl::math::dense > 
cpl::math::triangular_solve( 
  cpl::math::matrix< T , tag1 > const& A , 
  cpl::math::matrix< T , tag2 > const& b
) {


  assert( A.rows() == A.columns() ) ;

  const long n = A.rows() ;
  assert( b.rows() == n ) ;
  assert( b.columns() == 1 ) ;


  cpl::math::matrix< T , cpl::math::dense > x( n , 1 ) ;

  if( upper_bandwidth( A ) == 0 ) {

    // Lower triangular; do forward substitution.

    for( long i = 1 ; i <= n ; ++i ) {

      T sum = T() ;

      for( long j = cpl::math::first_column( A , i ) ; j < i ; ++j ) {

        sum -= A( i , j ) * x( j , 1 ) ;

      }

      sum += b( i , 1 ) ;

      if( A( i , i ) == T() ) {

        throw std::runtime_error( "matrix is singular" ) ;

      }

      x( i , 1 ) = sum / A( i , i ) ;

    }

    return x ;

  }

  if( lower_bandwidth( A ) == 0 ) {
    
    // Upper triangular; do backward substitution.

    for( long i = n ; i >= 1 ; --i ) {

      T sum = T() ;

      for( long j = cpl::math::last_column( A , i ) ; j > i ; --j ) {

        sum -= A( i , j ) * x( j , 1 ) ;

      }

      sum += b( i , 1 ) ;

      if( A( i , i ) == T() ) {

        throw std::runtime_error( "matrix singular" ) ;

      }

      x( i , 1 ) = sum / A( i , i ) ;

    }

    return x ;

  }

  throw std::runtime_error( "matrix not triangular" ) ;

}


template< typename T >
template< typename tag >
cpl::math::matrix< T , cpl::math::banded > const&
cpl::math::matrix< T , cpl::math::banded >::operator=( cpl::math::matrix< T , tag > const& rhs ) {

  if( 
       rhs.rows()    != rows() 
    || rhs.columns() != columns() 
    || lower_bandwidth( rhs ) > lower_bandwidth_
    || upper_bandwidth( rhs ) > upper_bandwidth_
    ) {

    cpl::math::matrix< T , cpl::math::banded > 
      ret( 
        rhs.rows() , 
        rhs.columns() , 
        lower_bandwidth( rhs ) , 
        upper_bandwidth( rhs ) 
      ) ;

    cpl::math::assign( ret , rhs ) ;
    this->swap( ret ) ;

  } 
  else { cpl::math::assign( *this , rhs ) ; }

  return *this ;

}


#endif // CPP_LIB_BANDED_MATRIX_H
